import pcl
import pcl.pcl_visualization
import numpy as np
from scipy.spatial import cKDTree
# from matplotlib import pyplot as pp
import cv2
import open3d as o3d

def transformSPH2XYZ(_row):
    row = _row.copy()
    a = np.deg2rad(row[:, 0])  # 水平转角
    b = np.deg2rad(row[:, 1])  # 竖直转角
    r = row[:, 2]  # 距离
    row[:, 0] = r * np.sin(b) * np.cos(a)
    row[:, 1] = -r * np.sin(b) * np.sin(a)
    row[:, 2] = r * np.cos(b)
    return row

def denoise_float32(array):
    kernel = np.array((
        [0.0625, 0.125, 0.0625],
        [0.125, 0.25, 0.125],
        [0.0625, 0.125, 0.0625]), dtype=np.float32)
    
    denoised = cv2.filter2D(array, -1, kernel)
    return denoised

def polar_filter(data, angular_resolution, filter2=None):
    width = int(360 / angular_resolution)
    height = int(180 / angular_resolution)
    range_image = np.zeros((width, height, 3), dtype=np.float32)
    count = np.zeros((width, height), dtype=np.int32)
    h, v = np.mod(np.int32(data[:, 0] / angular_resolution), width), np.mod(np.int32(data[:, 1] / angular_resolution), height)
    polar = set()    
    for i in range(data.shape[0]):
        range_image[h[i], v[i], :] = range_image[h[i], v[i], :] + data[i, 2:5]
        count[h[i], v[i]] = count[h[i], v[i]] + 1
        polar.add((h[i], v[i]))
    
    count_idx = count.nonzero()
    range_image[count_idx[0], count_idx[1], 0] = range_image[count_idx[0], count_idx[1], 0] / count[count_idx[0], count_idx[1]]
    range_image[count_idx[0], count_idx[1], 1] = range_image[count_idx[0], count_idx[1], 1] / count[count_idx[0], count_idx[1]]
    range_image[count_idx[0], count_idx[1], 2] = range_image[count_idx[0], count_idx[1], 2] / count[count_idx[0], count_idx[1]]
    
    if filter2 is not None:
        range_image = filter2(range_image)

    xyzi = np.zeros((len(polar), 8), dtype=np.float32)
    polar = np.array(list(polar))
    xyzi[:, 0:2] = polar * angular_resolution
    xyzi[:, 2:5] = range_image[polar[:, 0], polar[:, 1], :]
    # xyzi[:, 5] = range_image[polar[:, 0], polar[:, 1], 0]
    xyzi[:, 5:] = transformSPH2XYZ(xyzi[:, :3])
    return range_image, xyzi

def filter_outlier(cloud):
    outlier_filter = cloud.make_statistical_outlier_filter()
    outlier_filter.set_mean_k(100)
    outlier_filter.set_std_dev_mul_thresh(1)

    # outlier_filter = cloud.make_RadiusOutlierRemoval()
    # outlier_filter.set_MinNeighborsInRadius(1)
    # outlier_filter.set_radius_search(5)

    return outlier_filter.filter()

if __name__ == "__main__":
    data = np.float32(np.load("./x2.npy"))
    cloud = pcl.PointCloud(transformSPH2XYZ(data[:, 1:4]))
    viewer = pcl.pcl_visualization.CloudViewing()
    viewer.ShowMonochromeCloud(cloud)
    pcl.save(cloud, './raw.ply')
    ri, xyz = polar_filter(data[:, 1:6], 0.2)    
    points = xyz[:, 5:]
    cloud = pcl.PointCloud(points)
    cloud = filter_outlier(cloud)
    # distance, std_error = planar_statistics(cloud)    
    pcl.save(cloud, './filtered.ply')

    dri, dxyz = polar_filter(data[:, 1:6], 0.2, filter2=denoise_float32)
    points = dxyz[:, 5:]
    cloud = pcl.PointCloud(points)
    cloud = filter_outlier(cloud)    
    # distance, std_error = planar_statistics(cloud)    
    pcl.save(cloud, './filtered_denoised.ply')

    viewer = pcl.pcl_visualization.CloudViewing()
    viewer.ShowMonochromeCloud(cloud)
    flag = True
    while flag:
        flag != viewer.WasStopped()

    # pcd = o3d.geometry.PointCloud()
    # pcd.points = o3d.utility.Vector3dVector(xyz[:, 0:3])    
    # pcd.colors = o3d.utility.Vector3dVector(xyz[:, 3:6])
    
    # o3d.visualization.draw_geometries([pcd])